"""
JSON日志使用示例
展示如何在训练脚本中集成JSON日志系统
"""
import time
from utils import create_json_logger

# ============================================================================
# 示例1：基本使用
# ============================================================================

def example_basic_usage():
    """基本使用示例"""
    print("\n" + "=" * 80)
    print("示例1：基本使用")
    print("=" * 80)

    # 1. 创建JSON日志记录器
    json_logger = create_json_logger(
        log_dir='logs/unet',
        model_name='unet',
        fold_idx=0
    )

    # 2. 记录配置
    config = {
        'batch_size': 8,
        'learning_rate': 0.0001,
        'num_epochs': 120,
        'model': 'UNet'
    }
    json_logger.log_config(config)

    # 3. 模拟训练过程
    for epoch in range(3):
        # 模拟训练
        train_metrics = {
            'loss': 0.5 - epoch * 0.1,
            'iou': 0.6 + epoch * 0.05,
            'dice': 0.7 + epoch * 0.05
        }

        # 模拟验证
        val_metrics = {
            'loss': 0.4 - epoch * 0.08,
            'iou': 0.7 + epoch * 0.06,
            'dice': 0.8 + epoch * 0.05,
            'miou': 0.75 + epoch * 0.05,
            'pixel_accuracy': 0.99,
            'precision': 0.85 + epoch * 0.02,
            'recall': 0.80 + epoch * 0.03
        }

        # 记录epoch信息
        json_logger.log_epoch(
            epoch=epoch,
            train_metrics=train_metrics,
            val_metrics=val_metrics,
            lr=0.0001,
            epoch_time=120.5
        )

        print(f"Epoch {epoch}: Val mIoU={val_metrics['miou']:.4f}")

    # 4. 记录最佳模型
    json_logger.log_best_model(
        epoch=2,
        metrics=val_metrics,
        checkpoint_path='checkpoints/unet/fold_0/best.pth'
    )

    # 5. 记录训练总结
    json_logger.log_summary(
        total_epochs=3,
        total_time=361.5,
        best_miou=0.85
    )

    print(f"\n日志已保存到: {json_logger.get_log_path()}")


# ============================================================================
# 示例2：在实际训练循环中使用
# ============================================================================

def example_training_loop():
    """实际训练循环示例"""
    print("\n" + "=" * 80)
    print("示例2：实际训练循环")
    print("=" * 80)

    # 创建日志记录器
    json_logger = create_json_logger('logs/example', 'example_model', fold_idx=0)

    # 记录配置
    config = {
        'model': 'ExampleModel',
        'batch_size': 16,
        'learning_rate': 0.001,
        'optimizer': 'Adam',
        'scheduler': 'CosineAnnealingLR'
    }
    json_logger.log_config(config)

    # 训练参数
    num_epochs = 5
    best_miou = 0.0
    total_time = 0

    # 训练循环
    for epoch in range(num_epochs):
        epoch_start = time.time()

        # ===== 训练阶段 =====
        print(f"\nEpoch {epoch}/{num_epochs-1}")
        print("-" * 40)

        # 模拟训练（实际训练代码在这里）
        train_loss = 0.5 * (0.9 ** epoch)
        train_iou = 0.6 + epoch * 0.05
        train_dice = 0.7 + epoch * 0.04

        train_metrics = {
            'loss': train_loss,
            'iou': train_iou,
            'dice': train_dice
        }

        print(f"Train - Loss: {train_loss:.4f}, IoU: {train_iou:.4f}")

        # ===== 验证阶段 =====
        # 模拟验证（实际验证代码在这里）
        val_loss = 0.4 * (0.9 ** epoch)
        val_iou = 0.7 + epoch * 0.06
        val_dice = 0.8 + epoch * 0.04
        val_miou = 0.75 + epoch * 0.05
        val_pa = 0.99 + epoch * 0.001
        val_precision = 0.85 + epoch * 0.02
        val_recall = 0.80 + epoch * 0.03

        val_metrics = {
            'loss': val_loss,
            'iou': val_iou,
            'dice': val_dice,
            'miou': val_miou,
            'pixel_accuracy': val_pa,
            'precision': val_precision,
            'recall': val_recall
        }

        print(f"Val   - Loss: {val_loss:.4f}, mIoU: {val_miou:.4f}, Dice: {val_dice:.4f}")

        # 计算epoch时间
        epoch_time = time.time() - epoch_start
        total_time += epoch_time

        # 记录到JSON日志
        json_logger.log_epoch(
            epoch=epoch,
            train_metrics=train_metrics,
            val_metrics=val_metrics,
            lr=0.001 * (0.95 ** epoch),  # 模拟学习率衰减
            epoch_time=epoch_time
        )

        # 检查是否是最佳模型
        if val_miou > best_miou:
            best_miou = val_miou
            print(f"*** 新的最佳模型！mIoU: {best_miou:.4f} ***")

            # 记录最佳模型
            json_logger.log_best_model(
                epoch=epoch,
                metrics=val_metrics,
                checkpoint_path=f'checkpoints/example/fold_0/best.pth'
            )

    # 记录训练总结
    json_logger.log_summary(
        total_epochs=num_epochs,
        total_time=total_time,
        best_miou=best_miou
    )

    print(f"\n训练完成！")
    print(f"最佳mIoU: {best_miou:.4f}")
    print(f"总训练时间: {total_time:.2f}秒")
    print(f"日志文件: {json_logger.get_log_path()}")


# ============================================================================
# 示例3：读取和分析JSON日志
# ============================================================================

def example_read_and_analyze():
    """读取和分析JSON日志示例"""
    print("\n" + "=" * 80)
    print("示例3：读取和分析JSON日志")
    print("=" * 80)

    import json
    import os

    log_path = 'logs/example/example_model_fold_0.json'

    if not os.path.exists(log_path):
        print(f"日志文件不存在: {log_path}")
        print("请先运行示例2生成日志文件")
        return

    # 读取JSON日志
    with open(log_path, 'r', encoding='utf-8') as f:
        log_data = json.load(f)

    # 提取信息
    metadata = log_data['metadata']
    best_model = log_data['best_model']
    summary = log_data['summary']
    epochs = log_data['training']['epochs']

    # 打印信息
    print(f"\n模型: {metadata['model_name']}")
    print(f"Fold: {metadata['fold_idx']}")
    print(f"开始时间: {metadata['start_time']}")

    print(f"\n最佳模型:")
    print(f"  Epoch: {best_model['epoch']}")
    print(f"  mIoU: {best_model['metrics']['miou']:.4f}")
    print(f"  Dice: {best_model['metrics']['dice']:.4f}")
    print(f"  检查点: {best_model['checkpoint_path']}")

    print(f"\n训练总结:")
    print(f"  总Epoch数: {summary['total_epochs']}")
    print(f"  总时间: {summary['total_time']:.2f}秒")
    print(f"  最佳mIoU: {summary['best_miou']:.4f}")

    # 提取训练曲线数据
    print(f"\n训练曲线数据:")
    val_miou_list = [e['val']['miou'] for e in epochs]
    print(f"  验证集mIoU: {val_miou_list}")

    # 找到最佳epoch
    best_epoch_idx = val_miou_list.index(max(val_miou_list))
    print(f"  最佳Epoch: {best_epoch_idx} (mIoU={val_miou_list[best_epoch_idx]:.4f})")


# ============================================================================
# 主函数
# ============================================================================

def main():
    """运行所有示例"""
    print("\n" + "=" * 80)
    print("JSON日志系统使用示例")
    print("=" * 80)

    # 运行示例
    example_basic_usage()
    example_training_loop()
    example_read_and_analyze()

    print("\n" + "=" * 80)
    print("所有示例运行完成！")
    print("=" * 80)
    print("\n提示：")
    print("1. 查看生成的JSON日志文件:")
    print("   - logs/unet/unet_fold_0.json")
    print("   - logs/example/example_model_fold_0.json")
    print("\n2. 使用分析脚本:")
    print("   python analyze_log.py --log logs/example/example_model_fold_0.json")
    print("\n3. 生成训练曲线:")
    print("   python analyze_log.py --log logs/example/example_model_fold_0.json --plot")
    print()


if __name__ == '__main__':
    main()
